import torch
from pytorch3d.ops import knn_gather, knn_points

k=5
xyz = torch.randn([16,128,3])
B, N, C = xyz.shape
idxb = torch.arange(B).view(-1, 1)

dists_knn, idx_knn, neighbors_knn = knn_points(
    xyz, xyz, K=k,
    return_nn=True,
    return_sorted=False)  # [B, N, k]

idx_knn_ = idx_knn.view(B, -1) # [B, N * k]
neighbors = xyz[idxb, idx_knn_].view(B, N, k, C)
pt_raw = torch.unsqueeze(xyz, dim=2).expand_as(neighbors)
neighbor_vector = pt_raw - neighbors   # [B, N, k, C]
distance = torch.sqrt(torch.sum(neighbor_vector ** 2, dim=-1, keepdim=True)).squeeze()

knn_feat = knn_gather(xyz, idx_knn)

print( neighbors.shape, (neighbors-knn_feat).sum() )
print ( neighbors_knn.shape, (neighbors-neighbors_knn).sum() )
print( dists_knn.shape, (distance**2-dists_knn).sum() )